from typing import Optional
import numpy as np
from common.img_utils.img_shape import BandOrderType


def crop(
    nda_img: np.ndarray,
    x_start: int,
    x_size: int,
    y_start: int,
    y_size: int,
    pad_value: Optional[int] = None,
    band_order: BandOrderType = BandOrderType.GDAL,
):
    """Crop the input ndarray image.

    Parameters
    ----------
    nda_img : np.ndarray
        A 3-dimension image in np.ndarray or ndarray like form.
    x_start : int
        Start index for x axis (width). Left index.
    x_size : int
        Crop size for x axis (width).
    y_start : int
        Start index for y axis (height). Top index.
    y_size : int
        Crop size for y axis (height).
    pad_value : Optional[int], optional
        Padding size around crop image.
        By default None.
    band_order : BandOrderType, optional
        Band order for the `nda_img`, can be one of "gdal", "image",
        BandOrderType.GDAL, or BandOrderType.IMAGE.
        By default BandOrderType.GDAL

    Returns
    -------
    np.ndarray
        Cropped 3-dimension image.

    Raises
    ------
    ValueError
        Raised if the given `band_order` is not a known valid input.
    """
    nda_img = nda_img.copy()
    if pad_value is None:
        pad_value = 0
    if band_order == BandOrderType.GDAL:
        x_left, x_right = _check_1D_boundary(
            x_start - pad_value, x_start + x_size + pad_value, nda_img.shape[2]
        )
        y_left, y_right = _check_1D_boundary(
            y_start - pad_value, y_start + y_size + pad_value, nda_img.shape[1]
        )
        nda_img = nda_img[:, y_left:y_right, x_left:x_right]
    elif band_order == BandOrderType.IMAGE:
        x_left, x_right = _check_1D_boundary(
            x_start - pad_value, x_start + x_size + pad_value, nda_img.shape[1]
        )
        y_left, y_right = _check_1D_boundary(
            y_start - pad_value, y_start + y_size + pad_value, nda_img.shape[0]
        )
        nda_img = nda_img[y_left:y_right, x_left:x_right, :]
    else:
        raise ValueError(f"Unknown band_order {band_order}")
    return nda_img


def flip(
    nda_img: np.ndarray,
    horizon: bool = True,
    vertical: bool = True,
    band_order: BandOrderType = BandOrderType.GDAL,
):
    """Flip the input ndarray image horizontally/vertically.

    Parameters
    ----------
    nda_img : np.ndarray
        A 3-dimension image in np.ndarray or ndarray like form.
    horizon : bool, optional
        Whether to flip horizontally. By default True.
    vertical : bool, optional
        Whether to flip vertically. By default True
    band_order : BandOrderType, optional
        Band order for the `nda_img`, can be one of "gdal", "image",
        BandOrderType.GDAL, or BandOrderType.IMAGE.
        By default BandOrderType.GDAL.

    Returns
    -------
    np.ndarray
        Flipped 3-dimension image.

    Raises
    ------
    ValueError
        Raised if the given `band_order` is not a known valid input.
    """
    nda_img = nda_img.copy()
    if band_order == BandOrderType.GDAL:
        if horizon:
            if vertical:
                nda_img = np.flip(nda_img, (1, 2))
            else:
                nda_img = np.flip(nda_img, 2)
        else:
            if vertical:
                nda_img = np.flip(nda_img, 1)
    elif band_order == BandOrderType.IMAGE:
        if horizon:
            if vertical:
                nda_img = np.flip(nda_img, (0, 1))
            else:
                nda_img = np.flip(nda_img, 1)
        else:
            if vertical:
                nda_img = np.flip(nda_img, 0)
    else:
        raise ValueError(f"Unknown band_order {band_order}")
    return nda_img


def rotate(
    nda_img: np.ndarray,
    angle: float,
    reshape: bool = True,
    band_order: BandOrderType = BandOrderType.GDAL,
):
    """Rotate the input ndarray image by any angle.
    By default, rotate counterclockwise.
    In the case of a negative angle value, the direction of rotation is clockwise.

    Parameters
    ----------
    nda_img : np.ndarray
        A 3-dimension image in np.ndarray or ndarray like form.
    angle : float
        The rotation angle in degrees.
        By default, rotate counterclockwise.
        In the case of a negative value, the direction of rotation is clockwise.
    reshape : bool
        If reshape is true, the output shape is adapted so that
        the input array is contained completely in the output.
        By default True.
    band_order : BandOrderType, optional
        Band order for the `nda_img`, can be one of "gdal", "image",
        BandOrderType.GDAL, or BandOrderType.IMAGE.
        By default BandOrderType.GDAL.

    Returns
    -------
    np.ndarray
        Rotated 3-dimension image.

    Raises
    ------
    ValueError
        Raised if the given `band_order` is not a known valid input.
    """
    nda_img = nda_img.copy()
    from scipy import ndimage

    if band_order == BandOrderType.GDAL:
        nda_img = ndimage.rotate(nda_img, angle, axes=(1, 2), reshape=reshape)
    elif band_order == BandOrderType.IMAGE:
        nda_img = ndimage.rotate(nda_img, angle, axes=(0, 1), reshape=reshape)
    else:
        raise ValueError(f"Unknown band_order {band_order}")

    return nda_img


def rotate90(
    nda_img: np.ndarray,
    k: int,
    band_order: BandOrderType = BandOrderType.GDAL,
):
    """The rotate90() function is used to rotate an array by 90 degrees
    in the plane specified by axes.
    By default, rotate counterclockwise.
    In the case of a negative k value, the direction of rotation is clockwise.

    Parameters
    ----------
    nda_img : np.ndarray
        A 3-dimension image in np.ndarray or ndarray like form.
    k : int
        Number of times the array is rotated by 90 degrees.
        By default, rotate counterclockwise.
        In the case of a negative value, the direction of rotation is clockwise.
    band_order : BandOrderType, optional
        Band order for the `nda_img`, can be one of "gdal", "image",
        BandOrderType.GDAL, or BandOrderType.IMAGE.
        By default BandOrderType.GDAL.

    Returns
    -------
    np.ndarray
        Rotated 3-dimension image.

    Raises
    ------
    ValueError
        Raised if the given `band_order` is not a known valid input.
    """
    nda_img = nda_img.copy()
    if band_order == BandOrderType.GDAL:
        nda_img = np.rot90(nda_img, k=k, axes=(1, 2))
    elif band_order == BandOrderType.IMAGE:
        nda_img = np.rot90(nda_img, k=k, axes=(0, 1))
    else:
        raise ValueError(f"Unknown band_order {band_order}")
    return nda_img


def _check_1D_boundary(
    left: int, right: int, size: int, boundary_treatment: str = "shrink"
):
    """Checking left and right index in 1D boundary (0 - (size-1)).
    Return the modified left and right index with in the 1D boundary.
    In "shrink" mode, the distance between left and right may be shrinked.
    In "shift" mode, the distance between left and right will not shrinked.

    Parameters
    ----------
    left : int
        Left index to be checked.
    right : int
        Right index to be checked
    size : int
        Boundary size.
    boundary_treatment : str, optional
        How to treat the boundary condition when the moving block meets the end of rows/columns.
        Either "shrink" to shrink the block so to be within the valid zone but with smaller size then specified,
        or "shift" to shift the block to the valid zone and keep the size to be the same as specified.
        By default "shrink".

    Returns
    -------
    int, int
        Left, right after checking boundary.

    Raises
    ------
    ValueError
        Raised if the given `boundary_treatment` is not a known valid input.
    """
    if boundary_treatment not in ("shrink", "shift"):
        raise ValueError("Unknown boundary_treatment {}".format(boundary_treatment))

    if boundary_treatment == "shrink":
        if left < 0:
            left = 0
        if right > size:
            right = size
        return left, right
    elif boundary_treatment == "shift":
        if left - right > size:
            left = 0
            right = size
        elif left < 0:
            right = 0 - left + right
            left = 0
        elif right > size:
            left = left - right + size
            right = size
        return left, right
